package jp.ac.tohoku.ecei.cl.www.rite;

import java.io.*;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;

import jp.ac.tohoku.ecei.cl.www.base.*;

import gnu.trove.map.hash.*;

public class RITEEval {

    public static File goldFile = null;
    public static File sysFile = null;
    public static boolean unittest = false;

    public static enum SysFileFormat {RITEXML, TAB};
    public static SysFileFormat sysFileFormat = SysFileFormat.TAB;
    
    public static void main(String[] args) {
        processArguments(args);

        Evaluator labE = new Evaluator();
        Alphabet labelAlphabet = new Alphabet();

        // for unittest
        HashMap<String, Evaluator> utEvalHash = new HashMap<String, Evaluator>();
        
        if (sysFile == null) {
            System.err.println("Usage: $ java -jar rite2eval.jar [options] -g [gold-data] -s [system-output-file]");
            System.err.println("-t\tsystem-output-file is a tab-delimited or space-delimited format (default)");
            System.err.println("-x\tsystem-output-file is a rite xml format");
            System.err.println("Checking the system output format: $ java -jar rite2eval.jar -s [system-output-file]");
            System.exit(1);
        }

        try {
            RITEXML2RITEExample goldPipe = new RITEXML2RITEExample();
            
            ArrayList<RITEExample> goldExs = null;
            if (goldFile != null) {
                goldExs = goldPipe.parseXML(new FileInputStream(goldFile));
            }
            HashMap<String, RITEExample> sysExsHash = new HashMap<String, RITEExample>();

            // read sysout
            if (sysFileFormat == SysFileFormat.RITEXML) {
                RITEXML2RITEExample sysPipe = new RITEXML2RITEExample(); // provisional
                ArrayList<RITEExample> sysExs = sysPipe.parseXML(new FileInputStream(sysFile));
                for (int i = 0; i < sysExs.size(); i++) {
                    RITEExample sysEx = sysExs.get(i);
                    sysExsHash.put(sysEx.getId(), sysEx);
                }
            } else {
                TabbedTextParser sysPipe = new TabbedTextParser();
                ArrayList<String[]> parsed = sysPipe.parse(new FileInputStream(sysFile));
                ArrayList<RITEExample> sysExs = new ArrayList<RITEExample>();
                for (int i = 0; i < parsed.size(); i++) {
                    String[] result = parsed.get(i);
                    String id = result[0];
                    String label = result[1];
                    double score = Double.parseDouble(result[2]);
                    RITEExample sysEx = new RITEExample(id, null, null, label, 0, "");
                    sysExsHash.put(sysEx.getId(), sysEx);
                }
            }

            if (goldFile != null && (goldExs.size() != sysExsHash.size())) {
                System.err.println("the number of examples are different between the gold data and the system output.");
                System.err.println("gold: "+goldExs.size()+"  sys: "+sysExsHash.size());
            }
            
            if (goldFile != null) {
                for (int i = 0; i < goldExs.size(); i++) {
                    RITEExample goldEx = goldExs.get(i);
                    String category = goldEx.getCategory();

                    RITEExample sysEx = null;
                    if (sysExsHash.containsKey(goldEx.getId())) {
                        sysEx = sysExsHash.get(goldEx.getId());
                    }
                    // add semantic label results
                    if (sysEx != null) {
                        labE.add(labelAlphabet.lookupIndex(goldEx.getLabel(), true), labelAlphabet.lookupIndex(sysEx.getLabel(), true));
                    } else {
                        labE.addTarget(labelAlphabet.lookupIndex(goldEx.getLabel(), true));
                    }

                    if (category != null) {
                        unittest = true;
                        if (!utEvalHash.containsKey(category)) {
                            utEvalHash.put(category, new Evaluator());
                        }
                        Evaluator e = utEvalHash.get(category);
                        if (sysEx != null) {
                            e.add(labelAlphabet.lookupIndex(goldEx.getLabel(), true), labelAlphabet.lookupIndex(sysEx.getLabel(), true));  
                        } else {
                            e.addTarget(labelAlphabet.lookupIndex(goldEx.getLabel(), true));
                        }
                    }
                }
            } else if (goldFile == null && sysFile != null) {
                System.err.println("The format of this file is valid.");
            }
        } catch (Exception e) {
            e.printStackTrace();
        }

        if (goldFile != null && sysFile != null) {
            int[] labels = labE.getLabels();

            System.out.print("------------------------------------------------------------\n");
            System.out.printf("|%5s|%5s|%19s|%19s|%6s|\n", "Label", "#", "Precision", "Recall", "F1");
            for (int i = 0; i < labels.length; i++) {
                System.out.print("|");
                System.out.printf("%5s|", labelAlphabet.lookupObject(labels[i]));
                System.out.printf("%5d|", labE.getNumberOfTarget(labels[i]));
                System.out.printf("%6.2f(%5d/%5d)|", labE.getPrecision(labels[i]), labE.getNumberOfCorrect(labels[i]), labE.getNumberOfReturned(labels[i]));
                System.out.printf("%6.2f(%5d/%5d)|", labE.getRecall(labels[i]), labE.getNumberOfCorrect(labels[i]), labE.getNumberOfTarget(labels[i]));
                System.out.printf("%6.2f|", labE.getF(labels[i]));
                System.out.print("\n");
            }
            System.out.print("------------------------------------------------------------\n");

            System.out.print("Accuracy:\t");
            System.out.printf("%3.2f(%6d/%6d)\n", labE.getTotalPrecision(), labE.getNumberOfCorrect(), labE.getNumberOfTarget());

            System.out.print("Macro F1:\t");
            System.out.printf("%3.2f \n", labE.getTotalMacroF());
            System.out.print("\n");

            // print confusion matrix
            System.out.println("Confusion Matrix");
            System.out.print("------------");
            for (int i = 0; i < labels.length; i++) {
                System.out.printf("%4s", "----");
            }
            System.out.print("-\n");

            System.out.printf("|gold \\ sys|", "");            
            for (int i = 0; i < labels.length; i++) {
                System.out.printf("%4s", labelAlphabet.lookupObject(labels[i]));
            }
            System.out.print("|\n");
            System.out.print("------------");
            for (int i = 0; i < labels.length; i++) {
                System.out.printf("%4s", "----");
            }
            System.out.print("-\n");

            for (int i = 0; i < labels.length; i++) {
                System.out.print("|");
                System.out.printf("%10s|", labelAlphabet.lookupObject(labels[i]));
                for (int j = 0; j < labels.length; j++) {
                    System.out.printf("%4d", labE.getNum(labels[i], labels[j]));
                }
                System.out.print("|\n");
            }
            System.out.print("------------");
            for (int j = 0; j < labels.length; j++) {
                System.out.printf("%4s", "----");
            }
            System.out.print("-\n");

            // print detailed results (for unittest)
            if (unittest) {
                System.out.print("\n");
                int yLabel = labelAlphabet.lookupIndex("Y", false);
                int nLabel = labelAlphabet.lookupIndex("N", false);
                System.out.print("-------------------------------------------------------------------------------------------------\n");
                System.out.printf("|%17s|%15s|%15s|%6s|%15s|%15s|%6s|\n", "Category", "Y:Prec", "Y:Rec", "Y:F1", "N:Prec", "N:Rec", "N:F1");
                System.out.print("-------------------------------------------------------------------------------------------------\n");                
                for (Iterator it = utEvalHash.keySet().iterator(); it.hasNext();) {
                    String cat = (String) it.next();
                    Evaluator e = utEvalHash.get(cat);
                    System.out.printf("|%17s|", cat);
                    System.out.printf("%6.2f(%3d/%3d)|", e.getPrecision(yLabel), e.getNumberOfCorrect(yLabel), e.getNumberOfReturned(yLabel));
                    System.out.printf("%6.2f(%3d/%3d)|", e.getRecall(yLabel), e.getNumberOfCorrect(yLabel), e.getNumberOfTarget(yLabel));
                    System.out.printf("%6.2f|", e.getF(yLabel));
                    System.out.printf("%6.2f(%3d/%3d)|", e.getPrecision(nLabel), e.getNumberOfCorrect(nLabel), e.getNumberOfReturned(nLabel));
                    System.out.printf("%6.2f(%3d/%3d)|", e.getRecall(nLabel), e.getNumberOfCorrect(nLabel), e.getNumberOfTarget(nLabel));
                    System.out.printf("%6.2f|", e.getF(nLabel));
                    //System.out.printf("%6.2f|", e.getTotalMacroF());
                    System.out.print("\n");
                }
                System.out.print("-------------------------------------------------------------------------------------------------\n");                
            }
        }
    }

    private static void processArguments(String[] args) {
		int idx = 0;
		while (idx < args.length) {
            if (args[idx].equals("-g")) {
				idx++;
				goldFile = new File(args[idx]);
				idx++;
            } else if (args[idx].equals("-s")) {
				idx++;
				sysFile = new File(args[idx]);
				idx++;
			} else if (args[idx].equals("-x")) {
                idx++;
                sysFileFormat = SysFileFormat.RITEXML;
			} else if (args[idx].equals("-t")) {
                idx++;
                sysFileFormat = SysFileFormat.TAB;
            } else {
				idx++;
			}
		}
    }
}
